import cv2
import SVM_ocr
import numpy as np
import matplotlib.pyplot as plt
from PyQt5.QtWidgets import *
from PyQt5 import QtGui
from LPR import Ui_Form
import sys,os,time


class myMainWindow(Ui_Form,QMainWindow):
    def __init__(self):
        super(myMainWindow, self).__init__()
        self.setupUi(self)
        self.PB1.clicked.connect(self.openimage)  #绑定按键1：打开文件
        self.PB2.clicked.connect(self.detect)      #绑定按键2：车牌检测
        self.PB4.clicked.connect(self.closeEvent)      #绑定按键3;退出键

    def openimage(self): #打开文件
        self.imgName, imgType = QFileDialog.getOpenFileName(self, "打开图片", "./test/", "*.jpg;;*.png;;All Files(*)")
        print(self.imgName)
        jpg = QtGui.QPixmap(self.imgName).scaled(self.label.width(), self.label.height())
        self.label.setPixmap(jpg)

    def closeEvent(self, event): #关闭程序
        result = QMessageBox.question(self, "Xpath Robot", "是否退出?",
                                                QMessageBox.Yes | QMessageBox.No)
        if (result == QMessageBox.Yes):
            event.accept()
        else:
            event.ignore()

    def display(self,window, img, width, block=False):  # 使用cv2来展示预测结果图片
        """显示打开的图片
        """
        cv2.namedWindow(window,
                        cv2.WINDOW_NORMAL)
        cv2.resizeWindow(window,
                         width,
                         (int)(width * img.shape[0] / img.shape[1]))
        cv2.imshow(window, img)
        if block:
            cv2.waitKey(0)
            cv2.destroyAllWindows()

    def close_op(self,img, size):
        """形态闭合操作
        """
        pre_element = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))  # 返回指定形状和尺寸的结构元素，第一个参数MORPH_RECT表示为矩形，
        # 第二个参数（2,2）表示内核尺寸，第三个参数表示锚点中心位置，默认为（-1，-1），表示锚点位于中心点
        ero_element = cv2.getStructuringElement(cv2.MORPH_RECT,
                                                (size, (int)(size / 8)))
        dil_element = cv2.getStructuringElement(cv2.MORPH_RECT,
                                                (size, (int)(size / 6)))
        pre = cv2.erode(img,
                        pre_element,
                        iterations=1)  # 腐蚀操作，理解为图像断开，裂缝变大
        dil = cv2.dilate(pre,
                         dil_element,
                         iterations=1)  # 膨胀操作，图像断开，裂缝变小
        ero = cv2.erode(dil,
                        ero_element,
                        iterations=1)  # 腐蚀操作
        return ero

    def close_op_y(self,img, size):
        ero_element = cv2.getStructuringElement(cv2.MORPH_RECT,
                                                ((int)(size / 8), size))
        dil_element = cv2.getStructuringElement(cv2.MORPH_RECT,
                                                ((int)(size / 6), size))
        dil = cv2.dilate(img,
                         dil_element,
                         iterations=1) #膨胀操作，图像断开，裂缝变小
        ero = cv2.erode(dil,
                        ero_element,
                        iterations=1) #腐蚀操作
        return ero

    def check_plate(self,img, rect):
        """检查车牌是否存在
        """
        return True

    def draw_rect(self,img, rect):
        """在图像中绘制矩形框区域
        """
        box = cv2.boxPoints(rect) #获取矩形框的四个顶点位置，浮点型
        box = np.intp(box)        #取整
        return cv2.polylines(img, [box],
                             isClosed=True,
                             color=(0, 0, 255),
                             thickness=1) #绘制任意形状的框

    def scale_rect(self,rect, scale=1.0):
        """缩放矩形区域
        """
        rect = list(rect)
        scale = scale * rect[1][0] - rect[1][0]
        rect[1] = [x + scale for x in rect[1]]
        return rect

    def getBlueChannel(self,img):
        """获取图像的蓝色和绿色通道
        """
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        lower = np.array([40, 88, 40])
        upper = np.array([124, 255, 255])
        mask = cv2.inRange(hsv, lower, upper)

        _, s, _ = cv2.split(hsv)
        img = cv2.bitwise_and(s, s, mask=mask)
        return img

    def adjust_box(self,img, box):
        """调整车牌区域
        """
        pts = np.array(np.where(img > 0)).T
        pts[:, [0, 1]] = pts[:, [1, 0]]
        new_box = []
        for corner in box:
            new_box.append(list(pts[np.argmin([np.linalg.norm(pt - corner) for pt in pts])]))
        new_box = np.array(new_box)
        new_box[0][0] -= 3
        new_box[3][0] -= 3
        return new_box

    def box_normalize(self,box):
        box = box[np.argsort(box[:, 1])]
        if box[0][0] > box[1][0]:
            box[[0, 1], :] = box[[1, 0], :]
        if box[2][0] < box[3][0]:
            box[[2, 3], :] = box[[3, 2], :]
        return box

    def extract_plate(self,src, rect):
        """函数功能：从矩形区域中提取车牌
        """
        rect = self.scale_rect(rect, 1.2)

        """裁切矩形区域"""
        img = np.zeros_like(src).astype('uint8')
        box = np.intp(cv2.boxPoints(rect))
        box = self.box_normalize(box)
        img = cv2.fillPoly(img, [box], (255, 255, 255))
        cv2.copyTo(src, img, img)

        """图像预处理操作"""
        img = self.getBlueChannel(img)
        img = cv2.GaussianBlur(img, (3, 3), 0, 0, cv2.BORDER_DEFAULT)
        th = (np.mean(img) + (np.max(img) - np.mean(img)) * 0.4)
        _, img = cv2.threshold(img, th, 255, cv2.THRESH_BINARY)
        img = self.close_op(img, 32)

        """调整车牌区域"""
        cv2.polylines(src, [box], isClosed=True, color=(0, 0, 255), thickness=1)
        box = self.adjust_box(img, box)
        # cv2.polylines(src, [box], isClosed=True, color=(0, 0, 255), thickness=1)

        """通过透视变换对车牌进行校正"""
        box = box.astype('float32')

        M = cv2.getPerspectiveTransform(box, np.float32([[0, 0], [179, 0], [179, 59], [0, 59]]))
        plate = cv2.warpPerspective(src, M, (180, 60))#这两行代码为透视变换：将成像投影到一个新的平面，对车牌进行校正

        return plate

    def plot_data(self,data):
        x = np.linspace(0, len(data), len(data))
        plt.plot(x, data)
        plt.show()

    def calc_potential(self,img):
        """计算势能
        """
        sum = np.sum(img, 0)
        pot = np.zeros_like(sum)
        p, rp = 0, 0
        for i in range(len(sum)):
            """Linearly increase exponential decay"""
            p = p + sum[i] if sum[i] else p / 5
            rp = rp + sum[len(sum) - i - 1] if sum[len(sum) - i - 1] else rp / 5
            pot[i] += p
            pot[len(sum) - i - 1] += rp
        # plot_data(pot)
        return pot

    def adjust_vision(self,img):
        """调整视野的大小以便更好的匹配
        """
        # display('ch', img, 100, block=True)
        sum = np.sum(img, 1)
        up, down = -1, -1
        for i in range(len(sum)):
            if sum[i] > 512:
                if up == -1:
                    up = i
                down = i
        height = down - up + 1
        img = img[up: down + 1, :]
        img = cv2.resize(img, ((int)(img.shape[1] * 1.2), height))
        padding = height - img.shape[1]
        left = (int)(padding / 2)
        right = (int)(padding - left)
        img = np.hstack((np.zeros((height, left)), img))
        img = np.hstack((img, np.zeros((height, right))))
        img = cv2.resize(img, (20, 20))
        return img

    def plate_split(self,img):
        """通过势能拆分字符
        """
        pot = self.calc_potential(img)
        limitpot = np.mean(pot) * 0.5
        bound = []
        left = 0
        for i in range(len(pot) - 1):
            if pot[i] < limitpot and pot[i + 1] >= limitpot:
                left = i
            elif pot[i] >= limitpot and pot[i + 1] < limitpot:
                bound.append([left, i])
                left = 0
        if left != 0:
            bound.append([left, len(pot - 1)])
        return bound

    def plate_cut_text(self,img):
        '''
        车牌字符分割，将分割好的车牌区域进行字符分割
        '''
        sum = np.sum(img, 1)
        limit = np.mean(sum) * 0.2
        bound = []
        start = 0
        for i in range(len(sum) - 1):
            if sum[i] < limit and sum[i + 1] >= limit:
                start = i
            elif sum[i] >= limit and sum[i + 1] < limit:
                bound.append([start, i])
                start = 0
        if start != 0:
            bound.append([start, len(sum - 1)])
        up, down = 0, 0
        for b in bound:
            if b[1] - b[0] > down - up:
                up = b[0]
                down = b[1]

        return img[up: down + 1, :]

    def plate_recognition(self,plate):  # 字符识别代码
        '''
        车牌识别核心代码，使用SVM对分割的车牌字符进行识别，并将识别结果返回
        '''
        reader = SVM_ocr.Reader()
        img = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)
        th = (np.mean(img) + (np.max(img) - np.mean(img)) * 0.2)
        _, img = cv2.threshold(img, th, 255, cv2.THRESH_BINARY)
        img = self.plate_cut_text(img)
        # display('plate th', img, 360)
        img3 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 进行色彩空间的转换
        _image1 = QtGui.QImage(img3[:], img3.shape[1], img3.shape[0], plate.shape[1] * 3,
                              QtGui.QImage.Format_RGB888)  # pyqt5转换成自己能放的图片格式
        jpg_out = QtGui.QPixmap(_image1).scaled(self.label_5.width(), self.label_5.height())  # 设置图片大小
        self.label_5.setPixmap(jpg_out)  # 将截取到的车牌图像显示到GUI界面中

        bound = self.plate_split(img)
        plate_res = []
        for b in bound:
            if b[1] - b[0] > 5 and b[1] - b[0] < 28:
                if len(plate_res) == 0:
                    ch = reader.recognize_chinese(self.adjust_vision(img[:, b[0]:b[1] + 1]))[0]  # 识别车牌汉字
                else:
                    ch = reader.recognize_alnum(self.adjust_vision(img[:, b[0]:b[1] + 1]))[0]  # 识别车牌字符
                plate_res.append(ch)  # 将识别结果添加到变量ch中
        return plate_res  # 返回车牌识别结果

    def detect(self):   #车牌检测函数
        # if len(sys.argv) > 1:
        #     src = cv2.imread('./test/' + sys.argv[1])
        # else:
        #     src = cv2.imread('./test/test1.jpg')
        src = cv2.imread(self.imgName) #读取检测图片
        """提取蓝色通道"""
        time_start = time.time() #记录检测开始时间
        hsv = cv2.cvtColor(src, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        lower = np.array([100, 90, 40])
        upper = np.array([124, 255, 255])
        mask = cv2.inRange(hsv, lower, upper)
        img = cv2.bitwise_and(s, s, mask=mask)

        """图像预处理"""
        img = cv2.equalizeHist(img)
        img = cv2.GaussianBlur(img, (3, 3), 0, 0, cv2.BORDER_DEFAULT)

        """检测图像中的纹理"""
        flipped = cv2.flip(img, 1)
        sobel1 = cv2.Sobel(img, cv2.CV_8U, 1, 0, ksize=3)
        sobel2 = cv2.flip(cv2.Sobel(flipped, cv2.CV_8U, 1, 0, ksize=3), 1)
        img = sobel1 / 2 + sobel2 / 2
        img = img.astype('uint8')

        """阈值设置"""
        th = (np.mean(img) + (np.max(img) - np.mean(img)) * 0.6)
        ret, img = cv2.threshold(img, th, 255, cv2.THRESH_BINARY)

        """形态学闭合操作"""
        img = self.close_op(img, 36)

        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < 200 or area > 50000:
                continue

            rect = cv2.minAreaRect(contour) #获取最小外接矩阵，中心点坐标，宽高，旋转角度
            if rect[1][0] * rect[1][1] < 666 or area < rect[1][0] * rect[1][1] * 0.6:
                continue

            dy, dx = contour.flatten().reshape(contour.shape[0], -1).T.ptp(1)
            cwb = rect[1][1] / rect[1][0] if rect[1][1] > rect[1][0] else rect[1][0] / rect[1][1]
            if dy < dx or cwb < 2.5 or cwb > 6:
                continue

            if not self.check_plate(src, rect):
                continue
            '''
            对提取到的车牌进行识别检测
            '''
            plate = self.extract_plate(src, rect) #plate是截取到的车牌图像
            img2 = cv2.cvtColor(plate, cv2.COLOR_BGR2RGB) #进行色彩空间的转换
            _image = QtGui.QImage(img2[:], img2.shape[1], img2.shape[0], plate.shape[1] * 3,
                                   QtGui.QImage.Format_RGB888)  # pyqt5转换成自己能放的图片格式
            jpg_out = QtGui.QPixmap(_image).scaled(self.label_3.width(), self.label_3.height())  # 设置图片大小
            self.label_3.setPixmap(jpg_out)  # 将截取到的车牌图像显示到GUI界面中

            plate_res = self.plate_recognition(plate) #将识别结果存入列表
            time_end = time.time()
            time_sum = time_end - time_start #计算检测时间
            print('%.2f'% time_sum)
            self.label_8.setText('%.2f'%time_sum + 'S') #将检测时间显示到GUI界面中，对结果保留两位小数
            plate_res = ''.join(str(i) for i in plate_res) #导出识别结果
            self.label_2.setText(plate_res) #将识别结果在界面中显示

if __name__ == '__main__':
    app = QApplication(sys.argv)
    vieo_gui = myMainWindow()
    vieo_gui.show()
    sys.exit(app.exec_())
